package com.demo.rag.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;


@RestController
@RequestMapping(value = "/ai")
public class RagController {


    private final ChatClient chatClient;

    private final VectorStore vectorStore;

    public RagController(ChatClient.Builder chatClientBuilder, VectorStore vectorStore) {
        this.vectorStore = vectorStore;
        this.chatClient = chatClientBuilder
                //系统提示词
                .defaultSystem("""
                        ##角色
                        你是一个Java架构师，请以友好的方式回复Java相关的问题。
                        今天的日期是 {current_date}
                        """)
                //默认的RAG增强提示
                .defaultAdvisors(QuestionAnswerAdvisor.builder(vectorStore)
                        .searchRequest(SearchRequest.builder().similarityThreshold(0.8d).topK(1).build())
                        .build())
                .build();
    }

    @GetMapping(value = "/content/get/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> contentGetStream(@RequestParam(value = "question", defaultValue = "你好") String question) {

//        // 获取上下文
//        List<Document> context = vectorStore.similaritySearch(question); // maxResults=1
//
//        // 构建增强提示
//        String enhancedPrompt = "已知信息：\n" + String.join("\n", context) +
//                               "\n\n问题：" + question + "\n\n回答：";

        return chatClient.prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .user(question)
                .system(item -> item.param("current_date", LocalDate.now().toString()))
                .stream()
                .content();
    }

}